﻿using System;
using System.Collections.Generic;
using System.IO;

namespace DNA
{
    public class DnaIndex
    {
        internal class Node
        {
            private Node _a = null;
            private Node _c = null;
            private Node _g = null;
            private Node _t = null;

            private string _leaf = null;
            private int _depth;

            public Node(int depth)
            {
                _depth = depth;
            }

            public Node()
                : this(0)
            {
            }


            internal bool Add(string polimer)
            {
                if (polimer.Length == _depth)
                {
                    _leaf = polimer;
                    return true;
                }

                char c = polimer[_depth]; 
                switch (c)
                {
                    case 'a':
                        _expand(ref _a,polimer, _depth);
                        break;
                    case 'c':
                        _expand(ref _c, polimer, _depth);
                        break;
                    case 'g':
                        _expand(ref _g, polimer, _depth);
                        break;
                    case 't':
                        _expand(ref _t, polimer, _depth);
                        break;
                    default:
                        throw new Exception(string.Format("Unexpected nucleotide symbol: {0}", c));
                }
                return true;
            }

            internal bool _expand(ref Node node,string polimer,int depth)
            {
                if (node == null)
                    node = new Node(depth+1);
                node.Add(polimer);
                return true;
            }

            internal bool Contains(string polimer)
            {
                if (polimer.Length < _depth)
                    return false;

                if (polimer.Length == _depth)
                    return _leaf == polimer;
                                  
                Node node = _charToNode(polimer[_depth]);
                if (node == null)
                    return false;

                return node.Contains(polimer);
            }

            internal Node _charToNode(char c)
            {
                switch (c)
                {
                    case 'a':
                        return _a;
                    case 'c':
                        return _c;
                    case 'g':
                        return _g;
                    case 't':
                        return _t;
                    default:
                        throw new Exception(string.Format("Unexpected nucleotide symbol: {0}", c));
                }
            }

            internal void findAllStartingWith(string start, List<string> found)
            {
                if (start.Length <= _depth)
                {
                    if (_leaf != null)
                        found.Add(_leaf);

                    if (_a != null)
                        _a.findAllStartingWith(start, found);

                    if (_c != null)
                        _c.findAllStartingWith(start, found);

                    if (_g != null)
                        _g.findAllStartingWith(start, found);

                    if (_t != null)
                        _t.findAllStartingWith(start, found);
                }
                else
                {
                    char c = start[_depth];
                    Node node = _charToNode(c);
                    if (node != null)
                        node.findAllStartingWith(start, found);
                }
            }

            internal void Remove(string polimer)
            {
                if (polimer.Length == _depth)
                {
                    _leaf = null;
                    return;
                }

                if (polimer.Length > _depth)
                {
                    char c = polimer[_depth];
                    Node node = _charToNode(c);
                    node.Remove(polimer);

                    if (node._a == null && node._c == null && node._g == null && node._t == null && node._leaf == null)
                    {
                        switch (c)
                        {
                            case 'a':
                                _a = null;
                                break;
                            case 'c':
                                _c = null;
                                break;
                            case 'g':
                                _g = null;
                                break;
                            case 't':
                                _t = null;
                                break;
                            default:
                                throw new Exception(string.Format("Unexpected nucleotide symbol: {0}", c));
                        }
                    }
                }
            }
        }

        private Node _root = null;

        public DnaIndex()
        {
            _root = new Node();
        }

        public DnaIndex(StringReader reader) : this()
        {
            string line = null;
            while ((line = reader.ReadLine()) != null)
                Add(line);
        }

        public DnaIndex(string data)
            : this(new StringReader(data))
        {
        }

        public DnaIndex(List<string> list) : this()
        {
            foreach (var nucl in list)
                Add(nucl);
        }

        public bool Contains(string polimer)
        {
            return _root.Contains(polimer);
        }

        public void Add(string polimer)
        {
            if (polimer.Length == 0)
                return;
            _root.Add(polimer.ToLower());
        }

        public List<string> findAllStartingWith(string start)
        {
            List<string> found = new List<string>();
            _root.findAllStartingWith(start, found);
            return found;
        }

        public List<string> getAll()
        {
            return findAllStartingWith("");
        }

        public void Remove(string polimer)
        {
            if (polimer.Length > 0)
                _root.Remove(polimer);
        }
    }
}